# coding=utf8
from typing import TypeVar, List, Optional
from pyjvm.parser import MethodDescriptor, MethodDescriptorParser
from pyjvm.classfile.clsfile import ClassFile
from pyjvm.classfile.constpool import *
from pyjvm.classfile.clsmember import *
from pyjvm.classfile.attrs import *
from pyjvm.classfile.clstype import *
from pyjvm.exceptions import NoSuchFieldError, NoSuchMethodError, IllegalAccessError
from pyjvm.utils import ClassNameUtil


class AccessFlags:
    ACC_PUBLIC = 0x0001  # class field method
    ACC_PRIVATE = 0x0002  # field method
    ACC_PROTECTED = 0x0004  # field method
    ACC_STATIC = 0x0008  # field method
    ACC_FINAL = 0x0010  # class field method
    ACC_SUPER = 0x0020  # class  用到invokespeical指令时，需要对父类方法做特殊处理
    ACC_SYNCHRONIZED = 0x0020  # method
    ACC_VOLATILE = 0x0040  # field
    ACC_BRIDGE = 0x0040  # method
    ACC_TRANSIENT = 0x0080  # field
    ACC_VARARGS = 0x0080  # method
    ACC_NATIVE = 0x0100  # method
    ACC_INTERFACE = 0x0200  # class
    ACC_ABSTRACT = 0x0400  # class       method
    ACC_STRICT = 0x0800  # method
    ACC_SYNTHETIC = 0x1000  # class field method
    ACC_ANNOTATION = 0x2000  # class
    ACC_ENUM = 0x4000  # class field


class RuntimeConstantPool(list):

    def __init__(self, klass, cp: ConstantPool):
        self.klass = klass
        const_len = len(cp)
        constants = [None] * const_len
        list.__init__(self, constants)
        i = 1
        while i < const_len:
            info = cp[i]
            if info.tag in ConstantInfoTag.ConstValSet:
                self[i] = info.value
                if info.tag in (ConstantInfoTag.Long, ConstantInfoTag.Double):
                    i += 1
            elif info.tag == ConstantInfoTag.String:
                self[i] = info.string
            elif info.tag == ConstantInfoTag.Class:
                self[i] = ClassRef(info, self)
            elif info.tag == ConstantInfoTag.FieldRef:
                self[i] = FieldRef(info, self)
            elif info.tag == ConstantInfoTag.MethodRef:
                self[i] = MethodRef(info, self)
            elif info.tag == ConstantInfoTag.InterfaceMethodRef:
                self[i] = InterfaceMethodRef(info, self)
            i += 1


TypeClass = TypeVar('TypeClass', bound='Class')


class Class:

    def __init__(self):
        self.name: Optional[str] = None
        self.package_name: Optional[str] = None
        self.super_class_name: Optional[str] = None
        self.interface_names: List[str] = []
        self.runtime_const_pool: Optional[RuntimeConstantPool] = None
        self.fields: List[Field] = []
        self.methods: List[Method] = []
        self.attributes: List[AttributeInfo] = []
        self.access_flags = 0

        self.loader = None
        self.super_class = None
        self.interfaces: List = []
        self.instance_slot_count: int = 0
        self.static_slot_count: int = 0
        self.static_vars = {}

        self.is_public = False
        self.is_private = False
        self.is_protected = False
        self.is_static = False
        self.is_interface = False
        self.is_abstract = False
        self.inited = False
        self.is_array_class = False
        self.is_primitive = False
        self.is_synthetic = False
        self.j_class = None
        self.__source_file = None

    def __str__(self):
        return 'Class{%s}' % (self.get_java_name())

    @staticmethod
    def new_class(cls_file: ClassFile):
        self = Class()
        self.name = cls_file.class_name
        self.package_name = cls_file.package_name
        self.super_class_name = cls_file.super_class_name
        self.interface_names = cls_file.interface_names
        self.runtime_const_pool = RuntimeConstantPool(self, cls_file.const_pool)
        self.fields = Field.get_fields(self, cls_file.fields)
        self.methods = Method.get_methods(self, cls_file.methods)
        self.attributes = cls_file.attributes
        access_flags = uint16(cls_file.access_flags)
        self.__init_access_flags(access_flags)
        self.loader = None
        self.super_class = None
        self.interfaces = [None] * len(self.interface_names)

        self.instance_slot_count = 0
        self.static_slot_count = 0
        self.access_flags = access_flags
        return self

    @property
    def source_file(self):
        if self.__source_file is None:
            for attr in self.attributes:
                if isinstance(attr, SourceFileAttribute):
                    self.__source_file = attr.source_file
                    break
            else:
                self.__source_file = ''
        return self.__source_file

    def __init_access_flags(self, flags):
        self.is_public = bool(flags & AccessFlags.ACC_PUBLIC)
        self.is_private = bool(flags & AccessFlags.ACC_PRIVATE)
        self.is_protected = bool(flags & AccessFlags.ACC_PROTECTED)
        self.is_static = bool(flags & AccessFlags.ACC_STATIC)
        self.is_interface = bool(flags & AccessFlags.ACC_INTERFACE)
        self.is_abstract = bool(flags & AccessFlags.ACC_ABSTRACT)
        self.is_super = bool(flags & AccessFlags.ACC_SUPER)
        self.is_synthetic = bool(flags & AccessFlags.ACC_SYNTHETIC)

    def get_java_name(self):
        return self.name.replace('/', '.')

    def get_init_method(self):
        return self.get_static_method('<clinit>', '()V')

    def get_static_method(self, name: str, descriptor: str):
        return self.get_method(name, descriptor, True)

    def get_instance_method(self, name: str, descriptor: str):
        return self.get_method(name, descriptor, False)

    def get_constructors(self, public_only=True):
        init = '<init>'
        if public_only:
            return [m for m in self.methods if m.is_public and m.name == init]
        return [m for m in self.methods if m.name == init]

    def get_fields(self, public_only=True):
        return self.fields if not public_only else [fd for fd in self.fields if fd.is_public]

    def get_constructor(self, descriptor: str):
        return self.get_instance_method('<init>', descriptor)

    def get_main_method(self):
        return self.get_static_method('main', '([Ljava/lang/String;)V')

    def __get_member(self, name: str, descriptor: str, is_static: bool, is_method=True):
        klass = self
        while klass:
            members = klass.methods if is_method else klass.fields
            if not members:
                continue
            for member in members:
                if member.is_static == is_static:
                    if member.name == name and member.descriptor == descriptor:
                        return member
            klass = klass.super_class

    def get_method(self, name, descriptor, is_static):
        return self.__get_member(name, descriptor, is_static)

    def get_methods(self, public_only=True):
        return self.methods if not public_only else [m for m in self.methods if m.is_public]

    def get_static_field(self, name, descriptor):
        return self.get_field(name, descriptor, True)

    def get_field(self, name, descriptor, is_static):
        return self.__get_member(name, descriptor, is_static, False)

    def get_instance_field(self, name, descriptor):
        return self.get_field(name, descriptor, False)

    def accessible_to(self, other) -> bool:
        return self.is_public or self.package_name == other.package_name

    def is_sub_class_of(self, other: TypeClass) -> bool:
        klass = self.super_class
        while klass:
            if klass is other:
                return True
            klass = klass.super_class
        return False

    def is_super_class_of(self, other: TypeClass) -> bool:
        return other.is_sub_class_of(self)

    def is_sub_interface_of(self, interface: TypeClass) -> bool:
        for itf in self.interfaces:
            if itf is interface or itf.is_sub_interface_of(interface):
                return True
        return False

    def is_super_interface_of(self, interface: TypeClass) -> bool:
        return interface.is_sub_interface_of(self)

    def is_implements(self, interface: TypeClass) -> bool:
        klass = self
        while klass:
            for itf in klass.interfaces:
                if itf is interface or itf.is_sub_interface_of(interface):
                    return True
            klass = klass.super_class
        return False

    def accessible_from(self, other: TypeClass) -> bool:
        if self is other:
            return True

        if other.is_array_class:

            if self.is_array_class:
                other_component_class = self.loader.get_component_class(other)
                self_component_class = self.loader.get_component_class(self)
                return other_component_class is self_component_class or self_component_class.accessible_from(
                    other_component_class)

            else:
                if self.is_interface:
                    return self.name in ('java/lang/Cloneable', 'java/io/Serializable')
                else:
                    return self.name == 'java/lang/Object'

        else:

            if other.is_interface:
                if self.is_interface:
                    return self.is_super_interface_of(other)
                else:
                    return self.name == 'java/lang/Object'
            else:
                if self.is_interface:
                    return other.is_implements(self)
                else:
                    return other.is_sub_class_of(self)

    def get_component_array_class(self) -> TypeClass:
        return self.loader.get_component_array_class(self)

    def get_component_class(self) -> TypeClass:
        return self.loader.get_component_class(self)

    def put_static_field(self, field, val):
        self.static_vars[id(field)] = val

    def set_static_field(self, name: str, descriptor: str, ref):
        field = self.get_static_field(name, descriptor)
        self.static_vars[id(field)] = ref

    def get_static_val(self, field):
        return self.static_vars[id(field)]


class ClassMember:

    def __init__(self, klass: Class, class_member_info: ClassMemberInfo):
        self.klass = klass
        self.name = class_member_info.name
        self.descriptor = class_member_info.descriptor
        self.attributes = class_member_info.attributes
        self.is_long_or_double = self.descriptor in set('DJ')
        self.__init_access_flags(class_member_info.access_flags)
        self.__attributes = None

    @property
    def signature(self):
        attr = self.attr(SignatureAttribute)
        if attr is None:
            return
        return attr.signature

    @property
    def annotation_bytes(self):
        attr = self.attr(RuntimeVisibleAnnotationsAttribute)
        if attr is None:
            return b''
        return attr.code

    @property
    def annotation_default_bytes(self):
        attr = self.attr(AnnotationDefaultAttribute)
        if attr is None:
            return b''
        return attr.code

    def __prepare_attribute(self):
        if self.__attributes is None:
            attrs = {}
            for attr in self.attributes:
                key = attr.name + 'Attribute'
                if key in attrs:
                    raise Exception('attr key is exist')
                attrs[key] = attr
            self.__attributes = attrs

    def attr(self, attr_type):
        self.__prepare_attribute()
        attr = self.__attributes.get(attr_type.__name__)
        return attr

    def __init_access_flags(self, access_flags):
        self.access_flags = access_flags
        self.is_static = bool(access_flags & AccessFlags.ACC_STATIC)
        self.is_final = bool(access_flags & AccessFlags.ACC_FINAL)
        self.is_public = bool(access_flags & AccessFlags.ACC_PUBLIC)
        self.is_protected = bool(access_flags & AccessFlags.ACC_PROTECTED)
        self.is_private = bool(access_flags & AccessFlags.ACC_PRIVATE)
        self.is_native = bool(access_flags & AccessFlags.ACC_NATIVE)
        self.is_abstract = bool(access_flags & AccessFlags.ACC_ABSTRACT)

    def accessible_to(self, d: Class) -> bool:
        if self.is_public:  # 如果字段是public，任何类都可以访问
            return True
        c = self.klass
        if self.is_protected:  # 如果字段protected，自有子类和同一个包下的可以访问
            return d is c or d.is_sub_class_of(c) \
                   or d.package_name == c.package_name
        if not self.is_private:  # 如果是default，只有同个包下可以访问
            return d.package_name == c.package_name
        return d is c

    def equals_def(self, name: str, descriptor: str):
        return self.descriptor == descriptor and self.name == name


class Field(ClassMember):

    def __init__(self, klass: Class, class_member_info: ClassMemberInfo):
        ClassMember.__init__(self, klass, class_member_info)

    @property
    def const_val(self):
        attr = self.attr(ConstantValueAttribute)
        if attr is None:
            return
        return attr.const_value

    def get_type(self):
        return self.klass.loader.load_class(ClassNameUtil.to_class_name(self.descriptor))

    def __str__(self):
        return 'Field{%s.%s%s}' % (self.klass.get_java_name(), self.name, self.descriptor)

    @staticmethod
    def get_fields(klass: Class, cls_members):
        return [Field(klass, info) for info in cls_members]


class ExceptionHandler:

    def __init__(self, entry: ExceptionTableEntry, runtime_const_pool: RuntimeConstantPool):
        self.start_pc = entry.start_pc
        self.end_pc = entry.end_pc
        self.handler_pc = entry.handler_pc
        self.catch_type = ExceptionHandler.__get_catch_type(
            entry.catch_type_index, runtime_const_pool)

    @staticmethod
    def __get_catch_type(index: int, runtime_const_pool: RuntimeConstantPool):
        if index:
            return runtime_const_pool[index]


class ExceptionTable(list):

    def __init__(self, ex_entries: List[ExceptionTableEntry], runtime_const_pool: Optional[RuntimeConstantPool]):
        list.__init__(self, [ExceptionHandler(
            entry, runtime_const_pool) for entry in ex_entries])

    def find_ex_handler(self, ex_class: Class, pc: int) -> ExceptionHandler:
        for handler in self:
            if handler.start_pc <= pc < handler.end_pc:
                if handler.catch_type is None:
                    return handler
                catch_class: Class = handler.catch_type.resolved_class()
                if catch_class is ex_class or catch_class.is_sub_class_of(ex_class):
                    return handler


class Method(ClassMember):

    def __init__(self):
        self.klass: Optional[Class] = None
        self.name: Optional[str] = None
        self.descriptor: Optional[str] = None
        self.attributes: List[AttributeInfo] = []
        self.is_long_or_double = False
        self.max_stack = 0
        self.max_locals = 0
        self.code = b''
        self.exception_table = ExceptionTable([], None)
        self.__ex_types: List[Class] = []
        self.args_slot_count = 0
        self.parameter_types = []
        self.return_type = ''
        self.__parameters_classes = None
        self.__return_type_class = None

    @staticmethod
    def new(klass: Class, class_member_info: ClassMemberInfo):
        self = Method()
        ClassMember.__init__(self, klass, class_member_info)
        self.max_stack = 0
        self.max_locals = 0
        self.code = b''
        self.line_number_table = None
        self.__ex_types: List[Class] = []

        method_descriptor = MethodDescriptorParser(self.descriptor).parse()
        self.args_slot_count = Method.count_args_slot(method_descriptor)
        self.parameter_types = method_descriptor.parameter_types
        self.return_type = method_descriptor.return_type
        self.__init_code_attr()
        self.__init_native(method_descriptor)
        return self

    def get_parameter_types(self) -> List[Class]:
        if self.__parameters_classes is None:
            self.__parameters_classes = [self.klass.loader.load_class(ClassNameUtil.to_class_name(name))
                                         for name in self.parameter_types]
        return self.__parameters_classes

    def get_line_number(self, pc: int) -> int:
        if self.is_native:
            return -2
        if self.line_number_table is None:
            return -1
        return self.line_number_table.get_line_number(pc)

    def get_return_type(self) -> Class:
        if self.__return_type_class is None:
            self.__return_type_class = self.klass.loader.load_class(
                ClassNameUtil.to_class_name(self.return_type))
        return self.__return_type_class

    def __init_code_attr(self):
        attr = self.attr(CodeAttribute)
        if attr is None:
            return
        self.code_attr = attr
        self.max_stack = attr.max_stack
        self.max_locals = attr.max_locals
        self.code = attr.code
        self.line_number_table = attr.line_number_attr
        self.exception_table = ExceptionTable(
            attr.exception_table, self.klass.runtime_const_pool)

    @property
    def parameter_annotation_bytes(self):
        attr = self.attr(RuntimeVisibleParameterAnnotationsAttribute)
        if attr is None:
            return []
        return attr.code

    @property
    def exception_index_table(self):
        attr = self.attr(ExceptionsAttribute)
        if attr is None:
            return []
        return attr.exception_index_table

    def get_exception_types(self):
        if self.__ex_types is None:
            self.__ex_types = []
            for index in self.exception_index_table:
                class_ref = self.klass.runtime_const_pool[index]
                self.__ex_types.append(class_ref.resolved_class())
        return self.__ex_types

    def __init_native(self, md: MethodDescriptor):
        if self.is_native:
            self.__inject_code_attr(md.return_type)

    def __inject_code_attr(self, return_type: str):
        self.max_stack = 4
        self.max_locals = self.args_slot_count + (0 if self.is_static else 1)
        self.code = bytes({
                              'V': (0xfe, 0xb1),
                              'D': (0xfe, 0xaf),
                              'F': (0xfe, 0xae),
                              'J': (0xfe, 0xad),
                              'L': (0xfe, 0xb0),
                          }.get(return_type[0], (0xfe, 0xac)))

    def __str__(self):
        return 'Method{%s.%s%s}' % (self.klass.get_java_name(), self.name, self.descriptor)

    @staticmethod
    def get_methods(klass: Class, class_member_infos: list) -> list:
        return [Method.new(klass, info) for info in class_member_infos]

    @staticmethod
    def count_args_slot(md: MethodDescriptor) -> int:
        count_of_slot = 0
        for param_type in md.parameter_types:
            count_of_slot += 2 if param_type in set('DJ') else 1
        return count_of_slot

    def find_ex_handler_pc(self, ex_class: Class, pc: int):
        handler = self.exception_table.find_ex_handler(ex_class, pc)
        if handler is not None:
            return handler.handler_pc


class SymRef:
    """
    符号引用
    """

    def __init__(self, class_name: str, runtime_const_pool: RuntimeConstantPool):
        self.class_name = class_name
        self.runtime_const_pool = runtime_const_pool
        self.__class = None

    def resolved_class(self) -> Class:
        if self.__class is None:
            self.__class = self.resolve_class_ref()
        return self.__class

    def resolve_class_ref(self) -> Class:
        d = self.runtime_const_pool.klass
        c = d.loader.load_class(self.class_name)
        if not c.accessible_to(d):
            raise IllegalAccessError()
        return c


class ClassRef(SymRef):
    """
    类符号引用
    """

    def __init__(self, const_cls_info: ConstantClassInfo, runtime_const_pool: RuntimeConstantPool):
        SymRef.__init__(self, const_cls_info.class_name,
                        runtime_const_pool)


class MemberRef(SymRef):

    def __init__(self, ref_info: BaseConstantRefInfo, runtime_const_pool: RuntimeConstantPool):
        SymRef.__init__(self, ref_info.class_name, runtime_const_pool)
        self.name = ref_info.name
        self.descriptor = ref_info.descriptor


class FieldRef(MemberRef):

    def __init__(self, field_ref_info: ConstantFieldRefInfo, runtime_const_pool: RuntimeConstantPool):
        MemberRef.__init__(self, field_ref_info, runtime_const_pool)
        self.__field = None

    def resolved_field(self) -> Field:
        if self.__field is None:
            self.__field = self.resolve_field_ref()
        return self.__field

    def resolve_field_ref(self) -> Field:
        d = self.runtime_const_pool.klass
        c = self.resolved_class()
        field = lookup_field(c, self.name, self.descriptor)
        if field is None:
            raise NoSuchFieldError()
        if not field.accessible_to(d):
            raise IllegalAccessError()
        return field


class MethodRef(MemberRef):

    def __init__(self, ref_info: ConstantMethodRefInfo,
                 runtime_const_pool: RuntimeConstantPool):
        MemberRef.__init__(self, ref_info, runtime_const_pool)
        self.__method = None

    def resolved_method(self) -> Method:
        if self.__method is None:
            self.__method = self.resolve_method_ref()
        return self.__method

    def resolve_method_ref(self) -> Method:
        d = self.runtime_const_pool.klass
        c = self.resolved_class()
        method = lookup_method(c, self.name, self.descriptor)
        if method is None:
            raise NoSuchMethodError()
        if not method.accessible_to(d):
            raise IllegalAccessError()
        return method


class InterfaceMethodRef(MemberRef):

    def __init__(self, ref_info: ConstantInterfaceMethodRefInfo,
                 runtime_const_pool: RuntimeConstantPool):
        MemberRef.__init__(self, ref_info, runtime_const_pool)
        self.__method = None

    def resolved_method(self) -> Method:
        if self.__method is None:
            self.__method = self.__resolve_method_ref()
        return self.__method

    def __resolve_method_ref(self):
        d = self.runtime_const_pool.klass
        c = self.resolved_class()
        method = lookup_method_in_interface(c, self.name, self.descriptor)
        if method is None:
            raise NoSuchMethodError()
        if not method.accessible_to(d):
            raise IllegalAccessError()
        return method


def lookup_field(c: Class, name: str, descriptor: str) -> Field:
    for field in c.fields:
        if field.equals_def(name, descriptor):
            return field

    for interface in c.interfaces:
        field = lookup_field(interface, name, descriptor)
        if field:
            return field

    if c.super_class:
        return lookup_field(c.super_class, name, descriptor)


def lookup_method_in_methods(methods: list, name: str, descriptor: str) -> Method:
    for method in methods:
        if method.equals_def(name, descriptor):
            return method


def lookup_method_in_class(c: Class, name: str, descriptor: str) -> Method:
    klass = c
    while klass:
        method = lookup_method_in_methods(klass.methods, name, descriptor)
        if method:
            return method
        klass = klass.super_class


def lookup_method_in_interfaces(interfaces, name: str, descriptor: str) -> Method:
    for interface in interfaces:
        method = lookup_method_in_methods(interface.methods, name, descriptor)
        if method:
            return method
        method = lookup_method_in_interfaces(
            interface.interfaces, name, descriptor)
        if method:
            return method


def lookup_method_in_interface(interface: Class, name: str, descriptor: str) -> Method:
    method = lookup_method_in_methods(interface.methods, name, descriptor)
    if method is None:
        method = lookup_method_in_interfaces(
            interface.interfaces, name, descriptor)
    return method


def lookup_method(c: Class, name: str, descriptor: str) -> Method:
    method = lookup_method_in_class(c, name, descriptor)
    if method is None:
        method = lookup_method_in_interfaces(c.interfaces, name, descriptor)
    return method


__all__ = ['Class', 'RuntimeConstantPool', 'ClassMember', 'MemberRef',
           'FieldRef', 'Field', 'Method', 'MethodRef',
           'lookup_method_in_class', 'AccessFlags']
